{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcbe11cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import triton\n",
    "import torch\n",
    "import os\n",
    "import triton.language as tl\n",
    "import cutlass\n",
    "import cutlass.cute as cute\n",
    "from cutlass.cute.runtime import from_dlpack\n",
    "\n",
    "\n",
    "DEVICE = torch.device(\"cuda:0\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b98d5724",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we define the torch code\n",
    "\n",
    "def torch_add(x, y):\n",
    "    return x + y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d9562b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we define the torch compile code\n",
    "\n",
    "@torch.compile\n",
    "def torch_compile_add(x, y):\n",
    "    return x + y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1fff8d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we define the triton code\n",
    "\n",
    "@triton.jit\n",
    "def add_kernel(x_ptr,\n",
    "               y_ptr,\n",
    "               output_ptr,\n",
    "               n_elements,\n",
    "               BLOCK_SIZE: tl.constexpr):\n",
    "    pid = tl.program_id(axis=0)\n",
    "    # todo\n",
    "\n",
    "def triton_add(x, y):\n",
    "    output = torch.empty_like(x)\n",
    "    assert x.device == DEVICE and y.device == DEVICE and output.device == DEVICE\n",
    "    # todo\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "920701d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we define the cute code\n",
    "\n",
    "os.environ[\"PYTHONUNBUFFERED\"] = \"1\"\n",
    "\n",
    "@cute.kernel\n",
    "def vectorized_elementwise_add_kernel(\n",
    "    gA: cute.Tensor,\n",
    "    gB: cute.Tensor,\n",
    "    gC: cute.Tensor,\n",
    "):\n",
    "    # todo\n",
    "    pass\n",
    "\n",
    "@cute.jit\n",
    "def vectorized_elementwise_add(\n",
    "    mA: cute.Tensor,\n",
    "    mB: cute.Tensor,\n",
    "    mC: cute.Tensor\n",
    "):\n",
    "    # todo\n",
    "    pass\n",
    "\n",
    "a = torch.randn((2048, 2048), device=\"cuda\", dtype=torch.float32)\n",
    "b = torch.randn((2048, 2048), device=\"cuda\", dtype=torch.float32)\n",
    "c = torch.zeros((2048, 2048), device=\"cuda\", dtype=torch.float32)\n",
    "a_ = from_dlpack(a, assumed_align=16).mark_layout_dynamic()\n",
    "b_ = from_dlpack(b, assumed_align=16).mark_layout_dynamic()\n",
    "c_ = from_dlpack(c, assumed_align=16).mark_layout_dynamic()\n",
    "cute_dsl_add_compiled = cute.compile(vectorized_elementwise_add, a_, b_, c_)\n",
    "\n",
    "def cute_dsl_add(x, y):\n",
    "    out = torch.zeros_like(x)\n",
    "    cute_x = from_dlpack(x, assumed_align=16).mark_layout_dynamic()\n",
    "    cute_y = from_dlpack(y, assumed_align=16).mark_layout_dynamic()\n",
    "    cute_out = from_dlpack(out, assumed_align=16).mark_layout_dynamic()\n",
    "    cute_dsl_add_compiled(cute_x, cute_y, cute_out)\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f84fec3",
   "metadata": {},
   "outputs": [],
   "source": [
    "@triton.testing.perf_report(\n",
    "    triton.testing.Benchmark(\n",
    "        x_names=['size'],  # Argument names to use as an x-axis for the plot.\n",
    "        x_vals=[2**i for i in range(12, 28, 1)],  # Different possible values for `x_name`.\n",
    "        x_log=True,  # x axis is logarithmic.\n",
    "        line_arg='provider',  # Argument name whose value corresponds to a different line in the plot.\n",
    "        line_vals=['triton', 'torch', 'torch_compile', 'cute_dsl'],  # Possible values for `line_arg`.\n",
    "        line_names=['Triton', 'Torch', 'Torch Compile', 'Cute dsl'],  # Label name for the lines.\n",
    "        ylabel='GB/s',  # Label name for the y-axis.\n",
    "        plot_name='vector-add-performance',  # Name for the plot. Used also as a file name for saving the plot.\n",
    "        args={},  # Values for function arguments not in `x_names` and `y_name`.\n",
    "    ))\n",
    "def benchmark(size, provider):\n",
    "    x = torch.rand(size, device=DEVICE, dtype=torch.float32)\n",
    "    y = torch.rand(size, device=DEVICE, dtype=torch.float32)\n",
    "    quantiles = [0.5, 0.2, 0.8]\n",
    "    target_function = globals()[provider + \"_add\"]\n",
    "    z = target_function(x, y)\n",
    "    if torch.norm(z - (x + y)) > 1e-7 :\n",
    "        print(provider, torch.norm(z - (x + y)), x.shape)\n",
    "        print(z)\n",
    "        print(x + y)\n",
    "\n",
    "    ms, min_ms, max_ms = triton.testing.do_bench(lambda: target_function(x, y), quantiles=quantiles)\n",
    "    # gbps = lambda ms: 3 * x.numel() * x.element_size() * 1e-9 / (ms * 1e-3)\n",
    "    # todo for matrix add\n",
    "    return gbps(ms), gbps(max_ms), gbps(min_ms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c2f5e25",
   "metadata": {},
   "outputs": [],
   "source": [
    "benchmark.run(print_data=True, show_plots=True)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
